#define AnaSK_cxx
#include "AnaSK.h"
#include <TH2.h>
#include <TStyle.h>
#include <TCanvas.h>
#include <TRandom3.h>

#include "NeutrinoUtils.h"
#include "SKUtils.h"

void AnaSK::IniIO(const TString flist)
{
  printf("AnaSK::AnaSK flist %s\n", flist.Data());
  TChain * chain = NeutrinoUtils::InputFiles(flist, "h1");
  Init(chain);
}

//__________________________________________________________
// simulation related
//__________________________________________________________
void AnaSK::SetNeutrinoSim() const
{
  //check ipv[0] first
  //pnu(numnu)    :  particle code at neutrino interaction (P.D.G. code); numnu          :  number of particle at neutrino interaction
  //if(ipv[0]!= NeutrinoUtils::kGeantNeutrino || ipnu[0] != NeutrinoUtils::kPDGNumu){
  //if(ipv[0]!= NeutrinoUtils::kGeantNeutrino || ipnu[0] != NeutrinoUtils::kPDGAntiNumu){
  if(ipv[0]!= NeutrinoUtils::kGeantNeutrino){
    printf("AnaSK::SetNeutrinoSim wrong code for ipv[0] %d ipnu %d\n", ipv[0], ipnu[0]);
    printf("AnaSK::SetNeutrinoSim npar %d \n", npar);
    for(Int_t tmpii=0; tmpii<npar; tmpii++){
      printf("AnaSK::SetNeutrinoSim ipar %d ipv %d pmomv %f ipnu %d\n", tmpii, ipv[tmpii], pmomv[tmpii], ipnu[tmpii]);
    }
    exit(1); 
  }

  //TVector3 nup3(*NeutrinoUtils::kNeutrinoRec);
  //use real direction so that can check kNeutrinoRec = [0]-particle
  TVector3 nup3(dirv[0][0],dirv[0][1],dirv[0][2]);
  if(nup3.Mag()<EPSILON){
    printf("AnaSK::SetNeutrinoSim nup3 null!\n"); 
    nup3.Print();
    exit(1);
  }

  nup3 *= pmomv[0] * 1e-3 / nup3.Mag(); //in GeV/c

  const Double_t nuE = nup3.Mag();

  delete NeutrinoUtils::fNeutrinoSim;
  NeutrinoUtils::fNeutrinoSim = new TLorentzVector(nup3, nuE);
}

void AnaSK::SetMuonSim() const
{
  //
  //at simulation level
  //

  Int_t nfound = 0;

  //get muon: first lepton
//    numnu=1    incoming neutrino
//    :        2    target
//    :        3    outgoing lepton
//    :        4    outgoing target

  Int_t idmuon = NeutrinoUtils::kIniValue;
  //ipar=0 is beam neutrino
  for(Int_t ipar=1; ipar<npar; ipar++){
    if(ipv[ipar] == NeutrinoUtils::kGeantMuonMinus || ipv[ipar] == NeutrinoUtils::kGeantMuonPlus){
      if(idmuon==NeutrinoUtils::kIniValue){
        idmuon = ipar;
      }
      else{
        printf("AnaSK::SetMuonSim muonminus already exists!! nfound %d nrun %d nev %d idmuon %d ipv %d\n", nfound, nrun, nev, idmuon, ipv[idmuon]);
        for(Int_t tmpii=0; tmpii<npar; tmpii++){
          printf("AnaSK::SetMuonSim ipar %d ipv %d pmomv %f\n", tmpii, ipv[tmpii], pmomv[tmpii]);
        }
      } 
      nfound++;
    }
  }

  //sanity checks
  if(nfound){
    //sanity check 1
    if(idmuon!=2 && idmuon!=3){
      printf("AnaSK::SetMuonSim found muon not in expected index! should be NC, reset nfound to 0. nfound %d idmuon %d\n", nfound, idmuon);
      for(Int_t tmpii=0; tmpii<npar; tmpii++){
        printf("AnaSK::SetMuonSim test NC ipar %d ipv %d pmomv %f , ipnu[0,1,2,3,4] 0: %d, 1: %d, 2: %d, 3: %d, 4: %d\n", tmpii, ipv[tmpii], pmomv[tmpii], ipnu[0], ipnu[1], ipnu[2], ipnu[3], ipnu[4]);
      }
      nfound = 0;
    }
    //sanity check 2
    if(ipnu[0]==NeutrinoUtils::kPDGNumu){
      if(ipnu[idmuon]!= NeutrinoUtils::kPDGMuonMinus){
        printf("AnaSK::SetMuonSim wrong neutrino lepton code idmuon %d ipnu[0] %d ipnu[idmuon] %d\n", idmuon, ipnu[0], ipnu[idmuon]);
        for(Int_t tmpii=0; tmpii<npar; tmpii++){
          printf("AnaSK::SetMuonSim test lepton number ipar %d ipv %d pmomv %f , ipnu[0,1,2,3,4] 0: %d, 1: %d, 2: %d, 3: %d, 4: %d\n", tmpii, ipv[tmpii], pmomv[tmpii], ipnu[0], ipnu[1], ipnu[2], ipnu[3], ipnu[4]);
        }
        exit(1);
      }
    }
    else if(ipnu[0]==NeutrinoUtils::kPDGAntiNumu){
      if(ipnu[idmuon]!= NeutrinoUtils::kPDGMuonPlus){
        printf("AnaSK::SetMuonSim wrong antineutrino lepton code idmuon %d ipnu[0] %d ipnu[idmuon] %d\n", idmuon, ipnu[0], ipnu[idmuon]);
        for(Int_t tmpii=0; tmpii<npar; tmpii++){
          printf("AnaSK::SetMuonSim test lepton number ipar %d ipv %d pmomv %f , ipnu[0,1,2,3,4] 0: %d, 1: %d, 2: %d, 3: %d, 4: %d\n", tmpii, ipv[tmpii], pmomv[tmpii], ipnu[0], ipnu[1], ipnu[2], ipnu[3], ipnu[4]);
        }
        exit(1);
      }
    }
    else{
      printf("AnaSK::SetMuonSim wrong neutrino code %d\n", ipnu[0]); exit(1);
    }
  }
     

  //-------------------------------

  delete NeutrinoUtils::fMuonSim;
  NeutrinoUtils::fMuonSim = 0x0;

  //record first muon by definition
  //allow for multi-(+/-)muon 
  if(nfound >=1){
    TVector3 muonp3(dirv[idmuon][0], dirv[idmuon][1], dirv[idmuon][2]);
    if(muonp3.Mag()<EPSILON){
      printf("AnaSK::SetMuonSim muonp3 null!\n");
      muonp3.Print();
      exit(1);
    }

    muonp3 *= pmomv[idmuon] * 1e-3 / muonp3.Mag(); //in GeV/c
    
    const Double_t muE = TMath::Sqrt(NeutrinoUtils::MuonMass()*NeutrinoUtils::MuonMass()+muonp3.Mag2());
    //checked by ploting fMuonSim->M(), confirmed to be 0.1057
    NeutrinoUtils::fMuonSim = new TLorentzVector(muonp3, muE);
  }

  NeutrinoUtils::SetMuonNSim(nfound);
}

Bool_t AnaSK::IsAcceptQ2(const Bool_t kfast, TList *lin, const Int_t nsatu) const
{
  if(nsatu<0)
    return kTRUE;

  if(kfast){
    SetNeutrinoSim();
    SetMuonSim();
  }

  if(NeutrinoUtils::GetMuonNSim()==1){
    const Int_t nc = NeutrinoUtils::FillpreQ2(lin);
    if(nc<=nsatu){
      return kTRUE;
    }
  }
  
  return kFALSE;
}

//__________________________________________________________
// reconstructed
//__________________________________________________________
void AnaSK::SetNeutrinoRec() const
{
  delete NeutrinoUtils::fNeutrinoRec;

  //obtained from MC 14a via h1->Scan("dirv[0][0]:dirv[0][1]:dirv[0][2]","","colsize=30 precision=15")
  //new TVector3(0.669763624668121 , -0.742179155349731 , 0.0242228023707867);
  
  NeutrinoUtils::fNeutrinoRec = new TVector3(0.669763624668121 , -0.742179155349731 , 0.0242228023707867);
}

void AnaSK::SetNeutrinoGuess() const
{
 //
  //currect estimate
  //dirtot(3)      :  total direction sum ( each ring are weighted by RTOT )
  //
  delete NeutrinoUtils::fNeutrinoGuess;
  NeutrinoUtils::fNeutrinoGuess = 0x0;
  //test with truth
  //NeutrinoUtils::fNeutrinoGuess = new TVector3(*NeutrinoUtils::kNeutrinoRec);

  //+++
  //p.e.-weighted direction sum
  //NeutrinoUtils::fNeutrinoGuess = new TVector3(dirtot[0], dirtot[1], dirtot[2]);

  //+++

  //should only come after SetMuonRec
//  TVector3 psum;
//  for(Int_t iring=0; iring<NeutrinoUtils::GetMultiplicity(); iring++){
//    if(iring==NeutrinoUtils::fMuonID){
//      continue;
//    }
//    
//    TVector3 ivtmp(dir[iring][0], dir[iring][1], dir[iring][2]);
//    psum += ivtmp;
//  }
//  
//  psum *= 1./psum.Mag();
//  
//  NeutrinoUtils::fNeutrinoGuess = new TVector3(psum);

//+++

//default from the muon

  if(NeutrinoUtils::fMultiplicity>1){
    Int_t maxppid = NeutrinoUtils::fMuonID;
    Double_t mpp = -1e10;
    for(Int_t iring=0; iring<NeutrinoUtils::GetMultiplicity(); iring++){                                                                                                                                                      
      if(iring==NeutrinoUtils::fMuonID){
        continue;
      }
      
      const Double_t ipp = amomm[iring];
      if(ipp > mpp){
        maxppid = iring;
        mpp = ipp;
      }
    }

    const TVector3 pushdir(dir[maxppid][0], dir[maxppid][1], dir[maxppid][2]);
    const TVector3 pulldir(dirtot[0], dirtot[1], dirtot[2]);

    //about 10-15% below
    //const TVector3 pulldir(dir[NeutrinoUtils::fMuonID][0], dir[NeutrinoUtils::fMuonID][1], dir[NeutrinoUtils::fMuonID][2]);

    NeutrinoUtils::fNeutrinoGuess = new TVector3(pushdir+pulldir);
  }
}

void AnaSK::SetMuonRec() const
{
  //nring should be able to be 0
  if(nring>=10){
    printf("AnaSK::SetMuonRec nring %d >= size of NeutrinoUtils::fBetaRec\n", nring); exit(1);
  }
  NeutrinoUtils::SetMultiplicity(nring);

  Int_t nfound = 0;  

//  Int_t maxprobid = NeutrinoUtils::kIniValue;
//  Double_t mprob = -1e10;
//  
//  Int_t maxppid = NeutrinoUtils::kIniValue;
//  Double_t mpp = 0;

  Int_t maxpTid = NeutrinoUtils::kIniValue;
  Double_t mpT = 0;

  for(Int_t iring=0; iring<NeutrinoUtils::GetMultiplicity(); iring++){
    if(ip[iring] == SKUtils::kSKMuon){

      /*
      //maxprob ----------------------------------------------------
      //probms(6,nring):  probability of PID include opening angle  (approb(6,nring))
      //prmslg(6,nring):  probability of PID without opening angle (apprmslg(6,nring))
      //Float_t         probms[10][6];   //[nring]
      //Float_t         prmslg[10][6];   //[nring]  
      const Double_t iprob = probms[iring][SKUtils::kSKMuon];
      if(iprob > mprob){
        maxprobid = iring;
        mprob = iprob;
      }

      //maxpp -------------------------------------------------------
      const Double_t ipp = amomm[iring];
      if(ipp > mpp){
        maxppid = iring;
        mpp = ipp;
      }
      */

      //maxpT ----------------------------------------------------
      //true multi-muon event rare, rec. multi-muon due to misidentification of pion etc which is comming from soft pT(?)
      const Double_t ipt = NeutrinoUtils::GetPtRecFast(dir[iring][0],dir[iring][1],dir[iring][2],amomm[iring]);
      if(ipt > mpT){
        maxpTid = iring;
        mpT = ipt;
      }
      
      if(maxpTid>=0){
        nfound++;
      }
    }
  }

  //max pt efficiency best
  NeutrinoUtils::fMuonID = maxpTid; //maxprobid;

  //sanity check
  if(NeutrinoUtils::fMuonID<0 && nfound){
    printf("AnaSK::SetMuonRec() conflicting fMuonID %d nfound %d\n", NeutrinoUtils::fMuonID, nfound); exit(1);
  }

  delete NeutrinoUtils::fMuonRec;
  NeutrinoUtils::fMuonRec = 0x0;

  //multi-muon allow
  if(nfound>=1){
    TVector3 mup3(dir[NeutrinoUtils::fMuonID][0], dir[NeutrinoUtils::fMuonID][1], dir[NeutrinoUtils::fMuonID][2]);
    if(mup3.Mag()<EPSILON){
      printf("AnaSK::SetMuonRec mup3 null!\n");
      mup3.Print();
      exit(1);
    }

    mup3 *= amomm[NeutrinoUtils::fMuonID] * 1e-3 / mup3.Mag();
    
    const Double_t muE = TMath::Sqrt(NeutrinoUtils::MuonMass()*NeutrinoUtils::MuonMass()+mup3.Mag2());
    NeutrinoUtils::fMuonRec = new TLorentzVector(mup3, muE); 
  }

  NeutrinoUtils::SetMuonNRec(nfound);

  SetNeutrinoGuess();

  //checked -- OK: when maxpTdi = maxprobid, calcpt = maxpT
  //BTW, maxprob != maxpT, maxpp; but maxpT ~ maxpp
//  if(nfound>1){
//    static Int_t aa=0;
//    aa++;
//    printf("AnaSK::SetMuonRec nrun %d nev %d nfound %d maxprobid %d maxpTid %d maxppid %d maxpT %f calpt %f\n", nrun, nev, nfound, maxprobid, maxpTid, maxppid, mpT, NeutrinoUtils::GetMuonPtRec());
//    if(aa>10){
//      exit(1);
//    }
//  }
}

void AnaSK::SetBetaRec() const
{
  /*
//test
tree->Scan("fBetaMeanRec:fBetaRec[0]:fBetaRec[1]:fBetaRec[2]:fBetaRec[3]:fBetaRec[4]","fBetaRec[1]==-999&&fBetaMeanRec!=-999") //should be 0 entry
tree->Scan("fBetaMeanRec:fBetaRec[0]:fBetaRec[1]:fBetaRec[2]:fBetaRec[3]:fBetaRec[4]","fBetaRec[1]!=-999 && fBetaRec[2]==-999 && abs(fBetaMeanRec-(fBetaRec[0]+fBetaRec[1]-180))>1e-5")  //should be 0 entry; 1e-6 fail
tree->Scan("fBetaMeanRec:fBetaRec[0]:fBetaRec[1]:fBetaRec[2]:fBetaRec[3]:fBetaRec[4]","fMultiplicity==1 && fBetaMeanRec!=-999") //should be 0 entry
tree->Scan("fBetaMeanRec:fBetaRec[0]:fBetaRec[1]:fBetaRec[2]:fBetaRec[3]:fBetaRec[4]","fBetaRec[1]!=-999 && fBetaRec[2]!=-999 && fBetaRec[3]==-999 && abs(fBetaMeanRec*2-(fBetaRec[0]+fBetaRec[1]+fBetaRec[2]-180))>1e-5") ////should be 0 entry; 1e-6 fail    
  */

  for(Int_t iring=0; iring<NeutrinoUtils::GetMultiplicity(); iring++){
    NeutrinoUtils::fBetaRec[iring]= NeutrinoUtils::CalcBetaRec(dir[iring][0], dir[iring][1], dir[iring][2]);
    //printf("test %d %f\n", iring, NeutrinoUtils::fBetaRec[iring]);
  }


  NeutrinoUtils::fBetaRec[NeutrinoUtils::GetMultiplicity()]=-999;
}

void AnaSK::SetGammaRec() const
{
  /*
//test
tree->Scan("fGammaRec[0]:fGammaRec[1]:fGammaRec[2]:fGammaRec[3]:fGammaRec[fMuonID]:fAlphaRec","abs(fGammaRec[fMuonID]-fAlphaRec)>1e-10") //should be 0 entry
tree->Scan("fGammaRec[0]:fGammaRec[1]:fGammaRec[2]:fGammaRec[3]:fGammaRec[fMuonID]:fAlphaRec:fGammaMeanRec:fMultiplicity","fMultiplicity==1&&fGammaMeanRec!=-999") //should be 0 entry
tree->Scan("fGammaRec[0]:fGammaRec[1]:fGammaRec[2]:fGammaRec[3]:fGammaRec[fMuonID]:fAlphaRec:fGammaMeanRec:fMultiplicity","fMultiplicity>1&&fGammaMeanRec<0")  //should be 0 entry
tree->Scan("fGammaRec[0]:fGammaRec[1]:fGammaRec[2]:fGammaRec[3]:fGammaRec[fMuonID]:fAlphaRec:fGammaMeanRec:fMultiplicity","fMultiplicity==2&& abs(fGammaMeanRec - (fGammaRec[0]+fGammaRec[1]-fAlphaRec))>1E-10") //should be 0 entry
tree->Scan("fGammaMeanRec:fGammaRec[0]:fGammaRec[1]:fGammaRec[2]:fGammaRec[3]:fGammaRec[4]","fGammaRec[1]!=-999 && fGammaRec[2]!=-999 && fGammaRec[3]==-999 && abs(fGammaMeanRec*2-(fGammaRec[0]+fGammaRec[1]+fGammaRec[2]-fAlphaRec))>1e-5") ////should be 0 entry; 1e-6 fail    
  */

  for(Int_t iring=0; iring<NeutrinoUtils::GetMultiplicity(); iring++){
    NeutrinoUtils::fGammaRec[iring]= NeutrinoUtils::CalcGammaRec(dir[iring][0], dir[iring][1], dir[iring][2]);
  }

  NeutrinoUtils::fGammaRec[NeutrinoUtils::GetMultiplicity()] = -999;
}

//__________________________________________________________
//__________________________________________________________

void AnaSK::DoAnalysis(const TString filelist, const Int_t ntarget, const TString tag)
{
  IniIO(filelist);

  printf("AnaSK::DoAnalysis ntarget %d\n", ntarget);

  TList *lout=NeutrinoUtils::GetHistList();

  TTree *tout = NeutrinoUtils::GetTree();

  const Int_t nentries = fChain->GetEntriesFast();
  Int_t nsel=0;
  Int_t nfail = 0;

  TRandom3 oscRan(1);

  for (Int_t jentry=0; jentry<nentries;jentry++) {
    if(jentry%10000==0){
      printf("%d/%d nsel %d nfail %d\n", jentry, nentries,nsel, nfail);
    }

    if(!fChain->GetEntry(jentry)){
      printf("AnaSK::DoAnalysis GetEntry ends at jentry %d\n", jentry);
      break;
    }

    //--------------------------------------------
    NeutrinoUtils::FillCount(lout, 0);

    //save mode first
    NeutrinoUtils::fNeutMode = mode;
    NeutrinoUtils::fNeutrinoType = ipnu[0];  

    /*
    if(mode!=SKUtils::kSKCCDiffractive){
      nfail++;
      continue;
    }
    NeutrinoUtils::FillCount(lout, 1);
    */

    /*
    //it requires sim muon if nsatu>=0!
    if(!IsAcceptQ2(kTRUE, lout, -999)){
      nfail++;
      continue;
    }
    NeutrinoUtils::FillCount(lout, 3);
    */

    //considering oscillation
    NeutrinoUtils::fOsciMuMuProb = t2kposc[1][4];
    const Double_t rndm = oscRan.Rndm();
    if( rndm <NeutrinoUtils::fOsciMuMuProb){
      NeutrinoUtils::fkOsciMuSurvive = 1;
    }
    else{
      NeutrinoUtils::fkOsciMuSurvive = 0;
    }
    //---------------------------------------------------------------------------------------------------------
    //reconstruction
    //---------------------------------------------------------------------------------------------------------

    SetNeutrinoRec();
 
    SetMuonRec();

    /*
    //allow multiple muon
    if(NeutrinoUtils::GetMuonNRec()>1){
      nfail++;
      continue;
    }
    NeutrinoUtils::FillCount(lout, 4);
    */

    /*
    //allow 0-reconstructed in order to obtain efficiency
    if(NeutrinoUtils::GetMuonNRec()==0){
      nfail++;
      continue;
    }
    NeutrinoUtils::FillCount(lout, 5);
    */

    //---------

    NeutrinoUtils::SetMuonPtRec();

    SetBetaRec();
    SetGammaRec();

    NeutrinoUtils::SetRecKinematics();    

    //---------------------------------------------------------------------------------------------------------
    //simulation 
    //as of r2685, the histogram output are all about simulation (key words: neutrinoE, Q2, sim)
    //if run on real data, no hist will be filled 
    //tested by commenting out the following simulation part
    //---------------------------------------------------------------------------------------------------------

    SetNeutrinoSim();

    SetMuonSim();

    NeutrinoUtils::MuonNSimSanityCheck();

    /*
    //with single CC muon
    //open to allow multi-(+/-)muon real case: decay into muon and numu
    if(NeutrinoUtils::GetMuonNSim()>1){
      nfail++;
      continue;
    }
    NeutrinoUtils::FillCount(lout, 6);
    */

    /*
    //allow NC
    if(NeutrinoUtils::GetMuonNSim()==0){
      nfail++;
      continue;
    } 
    NeutrinoUtils::FillCount(lout, 7);
    */

    NeutrinoUtils::SetSimKinematics();    

    //---------------------------------------------------------------------------------------------------------
    //---------------------------------------------------------------------------------------------------------
    NeutrinoUtils::FillCount(lout,11);

    tout->Fill();
    
    NeutrinoUtils::DoFillTH(lout);

    nsel++;

    if(ntarget>=0 && nsel>=ntarget){
      break;
    }
  }

  printf("nentries %d nsel %d nfail %d delta %d\n", nentries, nsel, nfail, nentries-nsel-nfail);

  TFile *fout=new TFile(Form("SKAnaOutput_tag%s_%d_%d.root", tag.Data(), nsel, nfail),"recreate");

  tout->Write();
  //for further handling in plot
  lout->Write("lout",TObject::kSingleKey);

  fout->Save();
  fout->Close();
  delete fout;
}
